#在alphazero中使用MCTS搜索，利用策略-价值网络引导搜索

import numpy as np
import copy


class TreeNode(object):
    #决策树上一节点，每一节点跟踪自身价值Q，先验概率P,和由访问次数决定的先验得分U


    def __init__(self,parent,priorp):#初始化
        self.parent=parent
        self.child={}#(action(parent,prob),node)
        self.visit_time=0#访问次数
        self.value_Q=0#本节点价值
        self.prob_P=priorp#以此节点作为动作的概率
        self.score_u=0#以访问次数决定的得分

    def expand(self,action_prob):#扩展叶子节点，action_prob=[动作，动作概率]，概率由策略函数给出
        for action,prob in action_prob:
            if action not in self.child:
                self.child[action]=TreeNode(self,prob)

    def choose(self,cpuct):#在子节点中选择Q+U最大的节点，返回(动作，下一节点)
        return max(self.child.items(),
                    key=lambda actionnode: actionnode[1].get_value(cpuct))#赋值语句
    
    def update(self,leaf_value):#从叶子节点逆推当前节点价值，叶子节点价值是当前玩家视角的评估，有正负之分
        self.visit_time +=1
        self.value_Q += (leaf_value-self.value_Q)*1.0/self.visit_time#取相对于每次访问的评价价值

    def update_recursive(self,leaf_value):
        #若当前节点不是根节点，则先递归更新父节点
        if self.parent:
            self.parent.update_recursive(-leaf_value)
        self.update(leaf_value)

    def get_value(self,cpuct):#获得Q+u的值
        self.score_u=(cpuct*self.prob_P*np.sqrt(self.parent.visit_time)/(1+self.visit_time))
        return self.value_Q +self.score_u
    
    def Is_leaf(self):#判断是否是叶子节点
        if self.child=={}:
            return True
        else:
            return False
        
    def Is_root(self):#判断是否是根节点
        if self.parent:
            return False
        else:
            return True


class MCTS(object):
    #蒙特卡洛搜索

    def __init__(self,policy_value_function,cpuct=5,simu_time=10000):
        #policy_value对棋盘状态进行评估，进而输出一包含(动作，概率)的列表
        #以及在[-1,1]之间的得分(当前玩家视角的对对最终结果得分的期望值)
        #cpuct是一控制搜索到最大价值的速度的参数，其越大表明越依赖先前

        self.root=TreeNode(None,1.0)
        self.policy=policy_value_function
        self.cpuct=cpuct
        self.simu_time=simu_time#模拟次数

    def get_move_probs(self,state,temp=1e-3):
        #在模拟全部回合之后，得到所有可行动作和其概率，统计后得到动作概率
        #temp参数范围在(0,1]之间用于控制搜索

        for n in range(self.simu_time):#调用多次模拟
            state_cp=copy.deepcopy(state)
            self.simulation(state_cp)
        action_visit=[(action,node.visit_time) for action,node in self.root.child.items()] 
        acts,visits=zip(*action_visit)
        
        #通过归一化访问次数计算动作概率
        action_probs=softmax(1.0/temp*np.log10(np.array(visits)+1e-10))
        return acts,action_probs

    def simulation(self,state):
        #运行单次模拟，从根节点到叶子节点，获取最终结果价值，再反向更新其父节点价值和访问次数

        node=self.root
        while(1):
            if node.Is_leaf():
                break
            action,node=node.choose(self.cpuct)
            state.do_move(action)

        #利用策略函数计算当前叶子节点状态的价值，并输出(动作，概率)，以及得分
        action_probs,leaf_value=self.policy(state)#模拟后的状态state
        
        end,win=state.game_end()

        if not end:
            node.expand(action_probs)#未结束即还有节点未扩展
        else:
            if win==-1:
                leaf_value=0.0
            else:
                if win==state.get_current_player():leaf_value=1.0
                else:leaf_value=-1.0
        #更新模拟过程中经历节点的价值和访问次数
        node.update_recursive(-leaf_value)

    def update_move(self,last_move):
        #决策树前移，更新根节点

        if last_move in self.root.child:
            self.root=self.root.child[last_move]
            self.root.parent=None
        else:
            self.root=TreeNode(None,1.0)

class MCTSPlayer(object):
    #在MCTS算法基础上
    def __init__(self,policy_value_function,cpuct=5,simu_time=2000,selfplay=False):
        self.mcts=MCTS(policy_value_function,cpuct,simu_time)
        self.player=2
        self.selfplay=selfplay

    def set_player_ind(self, p):
        #两位玩家时需设置不同id，以便区分
        self.player = p

    def reset_player(self):
        #重置mcts根节点
        self.mcts.update_move(-1)


    def get_action(self,board,temp=1e-3,return_prob=0):
        avail_move=board.availables#获得可行动作
        move_probs=np.zeros(board.width*board.height)#矩阵初始化零

        #若有可行动作，获取动作概率，调用MCTS
        if len(avail_move)>0:
            acts,probs=self.mcts.get_move_probs(board,temp)
            move_probs[list(acts)]=probs                        
            if self.selfplay:
                if len(avail_move)==225:
                    move=112
                else:
                    #在自训练中需要加入dirichlet噪音
                    move=np.random.choice(acts,p=0.75*probs+0.25*np.random.dirichlet(0.03*np.ones(len(probs))))
                #更新树节点
                self.mcts.update_move(move)
            else:
                if len(avail_move)==225:
                    move=112
                else:
                    #比赛状态下使用最大概率的动作
                    move=np.random.choice(acts,p=probs)
                #更新树节点
                self.mcts.update_move(-1)
            if return_prob==1:
                return move, move_probs
            else:
                return move
        else:
            print('BOARD IS FULL')
            
    def __str__(self):
            return "MCTS {}".format(self.player)

def softmax(x):
    probs = np.exp(x - np.max(x))
    probs /= np.sum(probs)
    return probs
